import pandas as pd
import numpy as np
import logging
from .prophet_model import ProphetModel
from .lstm_model import LSTMModel
from .arima_model import ARIMAModel

logger = logging.getLogger(__name__)


class EnsembleModel:
    """集成多个模型的预测结果"""

    def __init__(self):
        self.prophet = ProphetModel()
        self.lstm = LSTMModel()
        self.arima = ARIMAModel()

        # 模型权重（可以根据历史准确度动态调整）
        self.weights = {
            'prophet': 0.4,
            'lstm': 0.35,
            'arima': 0.25
        }

    def predict(self, df, forecast_hours=24, confidence_level=0.95):
        """
        集成多个模型进行预测

        Args:
            df: DataFrame with columns ['timestamp', 'value']
            forecast_hours: 预测未来多少小时
            confidence_level: 置信区间水平

        Returns:
            预测结果列表
        """
        try:
            logger.info("开始集成模型预测...")

            predictions_dict = {}

            # 1. Prophet预测
            try:
                prophet_pred = self.prophet.predict(df, forecast_hours, confidence_level)
                predictions_dict['prophet'] = prophet_pred
                logger.info("Prophet预测成功")
            except Exception as e:
                logger.warning(f"Prophet预测失败: {str(e)}")

            # 2. LSTM预测（如果数据量足够）
            if len(df) >= 200:
                try:
                    lstm_pred = self.lstm.predict(df, forecast_hours, confidence_level)
                    predictions_dict['lstm'] = lstm_pred
                    logger.info("LSTM预测成功")
                except Exception as e:
                    logger.warning(f"LSTM预测失败: {str(e)}")

            # 3. ARIMA预测
            try:
                arima_pred = self.arima.predict(df, forecast_hours, confidence_level)
                predictions_dict['arima'] = arima_pred
                logger.info("ARIMA预测成功")
            except Exception as e:
                logger.warning(f"ARIMA预测失败: {str(e)}")

            # 如果所有模型都失败，抛出异常
            if not predictions_dict:
                raise Exception("所有预测模型都失败")

            # 4. 加权集成
            ensemble_results = self._weighted_ensemble(predictions_dict, forecast_hours)

            logger.info(f"集成预测完成，使用了{len(predictions_dict)}个模型")
            return ensemble_results

        except Exception as e:
            logger.error(f"集成预测失败: {str(e)}")
            raise

    def _weighted_ensemble(self, predictions_dict, forecast_hours):
        """加权平均集成多个模型的预测结果"""

        # 标准化权重
        available_models = list(predictions_dict.keys())
        total_weight = sum(self.weights[m] for m in available_models)
        normalized_weights = {
            m: self.weights[m] / total_weight
            for m in available_models
        }

        # 初始化结果数组
        ensemble_pred = np.zeros(forecast_hours)
        ensemble_lower = np.zeros(forecast_hours)
        ensemble_upper = np.zeros(forecast_hours)

        # 加权平均
        for model_name, predictions in predictions_dict.items():
            weight = normalized_weights[model_name]

            for i, pred in enumerate(predictions[:forecast_hours]):
                ensemble_pred[i] += pred['predicted_value'] * weight
                ensemble_lower[i] += pred['confidence_lower'] * weight
                ensemble_upper[i] += pred['confidence_upper'] * weight

        # 构建返回结果（使用第一个模型的时间戳）
        first_model_predictions = list(predictions_dict.values())[0]

        results = []
        for i in range(forecast_hours):
            results.append({
                'timestamp': first_model_predictions[i]['timestamp'],
                'predicted_value': max(0, int(ensemble_pred[i])),
                'confidence_lower': max(0, int(ensemble_lower[i])),
                'confidence_upper': max(0, int(ensemble_upper[i])),
                'ensemble_info': {
                    'models_used': available_models,
                    'weights': normalized_weights
                }
            })

        return results